from typing import ClassVar

from pandas import Index
from sklearn.feature_selection import SelectFromModel
from sklearn.linear_model import Ridge, Lasso
from sklearn.linear_model import ElasticNet
from sklearn.ensemble import RandomForestRegressor
from sklearn.ensemble import AdaBoostRegressor
from sklearn.ensemble import GradientBoostingRegressor
from selector._base_selector import _BaseSelector


class Embedded(_BaseSelector):
    """
    Embedded method of feature selection, subclass of _BaseFeatureSelector.
    Parameters:
        same with father class
    """

    def __init__(self, x, y, n_selected_features):
        super(Embedded, self).__init__(x, y, n_selected_features)

    def select_features(self, estimator: ClassVar) -> Index:
        """
        Parameters:
            estimator: ClassVar,
                {Ridge(), Lasso(), RandomForestRegressor(), AdaBoostRegressor(), GradientBoostingRegressor()}
        Return:
            Index, names of sorted selected features
        """
        selector = SelectFromModel(estimator=estimator, max_features=self.n_selected_features)
        selector.fit(self.x, self.y)
        if estimator == Ridge() or Lasso():
            selected_features = super(Embedded, self).sort_features(selector.estimator_.coef_)
        elif estimator == RandomForestRegressor() or AdaBoostRegressor() or GradientBoostingRegressor():
            selected_features = super(Embedded, self).sort_features(selector.estimator_.feature_importances_)
        else:
            raise NameError
        return selected_features
